#include <dim-sum/beehive.h>
#include <dim-sum/boot_allotter.h>
#include <dim-sum/cpu.h>
#include <dim-sum/fs.h>
#include <dim-sum/idle.h>
#include <dim-sum/irq.h>
#include <dim-sum/mem.h>
#include <dim-sum/sched.h>
#include <dim-sum/smp.h>
#include <dim-sum/string.h>
#include <dim-sum/stacktrace.h>
#include <dim-sum/syscall.h>
#include <dim-sum/timer.h>
#include <dim-sum/mm_types.h>
#include <dim-sum/exec.h>

#include <dim-sum/wait.h>
#include <kapi/dim-sum/task.h>

#include <asm/asm-offsets.h>
#include <asm/signal.h>

#include "internal.h"

static struct accurate_counter alloc_new_pid = {.counter = 0};

/**
 * 主CPU的idle线程堆栈
 * 使用静态分配的目的，是为了提供给CPU初始化过程使用
 */
__attribute__((aligned(PROCESS_STACK_SIZE))) 
union process_union master_idle_stack;

/**
 * 注意，是MAX_RT_PRIO + BITS_PER_LONG
 * 而不是MAX_RT_PRIO + BITS_PER_LONG **-1**
 * 这样最后总有一位是0
 */
#define SCHED_MASK_SIZE ((MAX_RT_PRIO + BITS_PER_LONG) / BITS_PER_LONG)
unsigned long sched_runqueue_mask[SCHED_MASK_SIZE] = { 0 };
struct double_list sched_runqueue_list[MAX_RT_PRIO+1];

union process_union *idle_proc_stacks[MAX_CPUS];
struct task_desc *idle_task_desc[MAX_CPUS];
#define INIT_SP(tsk)		((unsigned long)tsk + THREAD_START_SP)

struct smp_lock lock_all_task_list = SMP_LOCK_UNLOCKED(lock_all_task_list);
struct double_list sched_all_task_list = LIST_HEAD_INITIALIZER(sched_all_task_list);

unsigned int sched_timeslice = SCHED_TIMESLICE_TIME_DEFAULT;
/**
 * 分配/释放任务堆栈
 * 后期可以考虑使用对象池
 */
static void* alloc_proc_stack(void)
{
	return (void *)alloc_pages_memory(PAF_KERNEL, 
					get_order(PROCESS_STACK_SIZE));
}

void free_proc_stack(void* addr)
{
	free_pages_memory((unsigned long)addr, get_order(PROCESS_STACK_SIZE));
}

/**
 * 释放任务结构
 * 在任务的引用计数为0时调用
 */
void __release_task_desc(struct ref_count *ref)
{
	struct task_desc *tsk = container_of(ref, struct task_desc, ref);

	list_del_init(&tsk->all_list);
	free_proc_stack(tsk->stack);
}

static inline void add_to_runqueue(struct task_desc * p)
{
 	int pri = p->sched_prio;

	if (p->in_run_queue)
		return;

	list_insert_behind(&p->run_list,&(sched_runqueue_list[pri])); 
	p->in_run_queue = 1;
	if (p->sched_prio < current->sched_prio)
		set_task_need_resched(current);

	atomic_set_bit(pri, ( long unsigned int *)sched_runqueue_mask);  
}

static inline void del_from_runqueue(struct task_desc * p)
{
	int pri = p->sched_prio;

	if (!p->in_run_queue)
		return;

	list_del(&p->run_list);
	list_init(&p->run_list);

	if(list_is_empty(&(sched_runqueue_list[pri])))
	    atomic_clear_bit(pri, ( long unsigned int *)sched_runqueue_mask);

	p->in_run_queue = 0;
}

asmlinkage void __sched preempt_schedule(void)
{
	if (likely(preempt_count() || irqs_disabled()))
		return;
	
need_resched:
	/**
	 * 如果不添加PREEMPT_ACTIVE标志，会出现什么情况?
	 */
	add_preempt_count(PREEMPT_ACTIVE);
	schedule();
	sub_preempt_count(PREEMPT_ACTIVE);

	barrier();
	if (unlikely(test_process_flag(PROCFLAG_NEED_RESCHED)))
		goto need_resched;
}

asmlinkage void __sched preempt_in_irq(void)
{
	if (likely(preempt_count() || !irqs_disabled()))
		return;

need_resched:
	add_preempt_count(PREEMPT_ACTIVE);
	/**
	 * 这里不需要加内存屏障
	 */
	local_irq_enable();
	schedule();
	/**
	 * 在退回到中断汇编处理前，这里应当关闭中断
	 * 否则可能将堆栈击穿
	 */
	local_irq_disable();
	sub_preempt_count(PREEMPT_ACTIVE);

	barrier();
	if (unlikely(test_process_flag(PROCFLAG_NEED_RESCHED)))
		goto need_resched;
}

static struct task_desc *__switch_to(struct process_desc *prev,
	struct process_desc *new)
{
	struct task_desc *last;

	last = __switch_cpu_context(prev->task, new->task);

	return last;
}

asmlinkage void schedule(void)
{
	struct task_desc *prev, *next;
	unsigned long flags;
	int idx;

	if (irqs_disabled() || (preempt_count() & ~PREEMPT_ACTIVE)) {
		printk("cannt switch task, preempt count is %lx, irq %s.\n",
			preempt_count(), irqs_disabled() ? "disabled" : "enabled");
		BUG();
	}

need_resched:
	/**
	 * 这里必须手动增加抢占计数
	 * 以避免在打开锁的时候执行调度，那样就乱套了
	 */
	preempt_disable();
	smp_lock_irqsave(&lock_all_task_list, flags);

	prev = current;
	/**
	 * 很微妙的两个标志，请特别小心
	 */
	if (!(preempt_count() & PREEMPT_ACTIVE) && !(prev->state & TASK_RUNNING))
		del_from_runqueue(prev);

	idx = find_first_bit(sched_runqueue_mask, MAX_RT_PRIO + 1);
	/**
	 * 没有任务可运行
	 */
	if (idx > MAX_RT_PRIO)
		/**
		 * 选择本CPU上的IDLE任务来运行
		 */
		next = idle_task_desc[smp_processor_id()];
	else
		next = list_first_container(&sched_runqueue_list[idx],
							struct task_desc, run_list);

	/**
	 * 什么情况下，二者会相等??
	 */
	if (unlikely(prev == next)) {
		clear_task_need_resched(prev);
		smp_unlock_irq(&lock_all_task_list);
		preempt_enable_no_resched();
		goto same_process;
	}
	clear_task_need_resched(prev);

	next->prev_sched = prev;
	task_process_info(next)->cpu = task_process_info(prev)->cpu;
	prev = __switch_to(task_process_info(prev), task_process_info(next)); 
	barrier();
	/**
	 * 只能在切换后，释放上一个进程的结构。
	 * 待wait系统完成后，这里再修改
	 */
	if (current->prev_sched && (current->prev_sched->state & TASK_ZOMBIE)) {
		loosen_task_desc(current->prev_sched);
		current->prev_sched = NULL;
	}

	smp_unlock_irq(&lock_all_task_list);
	preempt_enable_no_resched();

same_process:
	if (unlikely(test_process_flag(PROCFLAG_NEED_RESCHED)))
		goto need_resched;

	return;
}

void __sched sched_update_timeslice(struct task_desc * p)
{
	if (p->sched_policy != SCHED_TIMESLICE || !(p->state & TASK_RUNNING))
		return;

	p->ts_time += TICK_NSEC / 1000000;
	if (p->ts_time >= sched_timeslice) {
		list_move_to_behind(&p->run_list,&(sched_runqueue_list[p->sched_prio]));
		set_task_need_resched(p);
		p->ts_time = 0;
	}
}

int __sched wake_up_process_special(struct task_desc *tsk, unsigned int state, int sync)
{
	unsigned long flags;
	int ret;

	smp_lock_irqsave(&lock_all_task_list, flags);
	if (tsk->state & state) {
		add_to_runqueue(tsk);
		tsk->state = TASK_RUNNING;
		ret = 0;
	} else
		ret = -1;
	smp_unlock_irqrestore(&lock_all_task_list, flags);

	return ret;
}

int __sched wake_up_process(struct task_desc *tsk)
{
	return wake_up_process_special(tsk, TASK_STOPPED | TASK_TRACED |
				 TASK_INTERRUPTIBLE | TASK_UNINTERRUPTIBLE, 0);
}

void change_task_prio(struct task_desc *tsk, int prio)
{
	unsigned long flags;
	
	if (prio == tsk->sched_prio)
		return;

	smp_lock_irqsave(&lock_all_task_list, flags);

	/**
	 * 此处需要用in_list进行判断，而不能判断进程状态是否为TASK_RUNNING
	 */
	if (tsk->in_run_queue) {
		del_from_runqueue(tsk);
		tsk->sched_prio = prio;
		add_to_runqueue(tsk);
	} else
		tsk->sched_prio = prio;

	smp_unlock_irqrestore(&lock_all_task_list, flags);
}

static int exit_task_sched(struct task_desc * tsk)
{
	if (!tsk)
		return -EINVAL;
	return 0;
}

static int exit_task_context(struct task_desc *tsk) 
{
	struct fp_context* f;
	if (!tsk)
		return -EINVAL;
	
	f = tsk->fp_regs;
	if (!f)
		return 0;

	smp_lock(&tsk->lock);
	tsk->fp_regs = NULL;
	smp_unlock(&tsk->lock);

	kfree(f);
	return 0;
}


static int exit_task(struct task_desc *tsk)
{
	if (!tsk)
		return -EINVAL;
	
	BUG_ON(exit_task_fs(tsk));
	BUG_ON(exit_task_sched(tsk));	
	BUG_ON(exit_task_context(tsk));
	BUG_ON(exit_execv_info(tsk));
	return 0;
}

static void __noreturn do_exit_task(int code)
{
	struct task_desc *tsk = current;
	struct process_desc *proc;

	preempt_disable();
	wake_up_interruptible(&tsk->wait_child_exit);

	set_current_state(TASK_ZOMBIE);
	
	tsk->exit_code = code;
	
	proc = task_on_process(tsk);
	proc->flags |= PF_EXITING;

	BUG_ON(proc->task != tsk);
	BUG_ON(exit_task(tsk));
	if (proc->mem)
		BUG_ON(exit_mem(proc, proc->mem));
	
	preempt_enable();
	/**
	 * 这里一旦切换出去以后，无法再回来
	 * 因此清理工作是在schedule中做的
	 * 虽然有点奇怪，但是以后会调整
	 */
	schedule();

	BUG();
}


/**
 * 任务最初的入口
 */
void task_entry(void)
{
	struct task_desc *tsk = current;
	int (*func)(void *) = tsk->task_main;
	int ret;

	/** 
	 * 稍微有点费解
	 * 这里是与schedule函数前半部分对应
	 */
	smp_unlock(&lock_all_task_list);
	local_irq_enable();
	preempt_enable();

	if (tsk->create_mode & TASK_TY_USER) {
		/**
		 * 获取argv, envp参数
		 */
		char* const *argv = (char* const *)arch_get_argv();
		char* const *envp = (char* const *)arch_get_envp();

		BUG_ON(vmem_execve((void*)tsk->bin_path, argv, envp));
		start_user_task(tsk);
		/**
		 * 多余的操作进程不会返回到这里
		 */
		do_exit_task(0);
	}
	/**
	 * 执行任务主函数
	 */
	ret = func(tsk->main_data);

	/**
	 * 任务执行完毕后，进行清理工作
	 */
	do_exit_task(ret);
}

asmlinkage __noreturn void sys_exit(int error_code)
{
	do_exit_task((error_code & 0xff)<<8);
}

/**
 * 分配一个新的PID
 */
pid_t get_new_pid(void)
{
	return accurate_inc(&alloc_new_pid);
}

static void setup_process_unoin(union process_union *stack,
		struct task_desc *tsk, int cpu, int preempt_count)
{
	if (stack == NULL || tsk == NULL)
		BUG();
	stack->process_desc.task = tsk;
	stack->process_desc.cpu = cpu;
	stack->process_desc.magic = TASK_MAGIC;
	stack->process_desc.preempt_count = preempt_count;
	stack->process_desc.flags = 0;
	stack->process_desc.current_irq = -1;

	/**
	 * 设置任务栈
	 */
	tsk->stack = stack;
}

static union process_union *init_task_process_union(struct task_desc * tsk)
{
	union process_union *stack = NULL;
	/* 创建进程描述符及堆栈 */
	stack = (union process_union *)alloc_proc_stack();
	if (stack == NULL)
		goto failed;

	/* 初始化任务process_desc, 对于非idle任务来说，一般都是属于创建自己的cpu。 */
	setup_process_unoin(stack, tsk, smp_processor_id(), INIT_TASK_PREEMPT_COUNT);

	/* 创建进程内存地址空间 */
	if (create_process_memspace(&stack->process_desc))
		goto failed;

	return stack;
failed:
	BUG();
	/**
	 * 理论上这里应该释放 create_process_memspace 申请的内存,
	 * 但之前没做,暂且这样。
	 */
	return NULL;
}

static inline void enqueue_task(struct task_desc *p)
{

}

static inline  void dequeue_task(struct task_desc *p)
{

}

static void init_task_desc(struct task_desc* tsk, pid_t pid)
{
	BUG_ON(!tsk);
	hold_task_desc(tsk);
	/* 对 task 默认赋值 */
	tsk->magic	= TASK_MAGIC;
	tsk->stack = NULL;
	smp_lock_init(&tsk->lock);
	tsk->state = TASK_INIT;
	tsk->flags = 0;
	tsk->atomic_flags = 0;
	tsk->exit_code = 0;
	tsk->in_run_queue = 0;
	tsk->timer_slack_ns = INIT_TIMER_SLACK_NS;
	tsk->pid = pid;
	tsk->in_syscall = 0;


	/**
	 * 任务组领头进程也指向自己
	 */
	tsk->gleader = tsk;
	tsk->addr_limit = TASK_SIZE;
	
	tsk->uid = 0;
	tsk->euid = 0;
	tsk->suid = 0;
	tsk->gid = 0;
	tsk->egid = 0;
	tsk->sgid = 0;
	tsk->fsuid = 0;
	tsk->fsgid = 0;

	tsk->parent_exec_id = 0;
	tsk->self_exec_id = 0;

	tsk->user_regs = NULL;

	tsk->tp_value = 0;
	tsk->set_child_tid = NULL;
	tsk->clear_child_tid = NULL;

	tsk->personality = 0;

	tsk->create_mode = 0;
	timer_init(&tsk->real_timer);
	// tsk->real_timer.handle = it_real_fn;
	tsk->real_timer.data = tsk;
	tsk->it_real_value = 0;
	tsk->it_real_incr = 0;
	/**
	 * 下面几个成员暂时未使用，默认置0
	 */
	tsk->min_flt = 0;
	tsk->maj_flt = 0;
	tsk->nswap = 0;
	tsk->cmin_flt = 0;
	tsk->cmaj_flt = 0;
	tsk->cnswap = 0;

	list_init(&tsk->run_list);
	list_init(&tsk->sched_all_list);
	list_init(&tsk->global_all_list);

	init_waitqueue(&tsk->wait_child_exit);
	return;
}

static inline struct task_desc *create_task_desc(void)
{
	struct task_desc *tsk;
	tsk = kzalloc(sizeof(struct task_desc), PAF_KERNEL);
	if (tsk == NULL)
		return NULL;

	init_task_desc(tsk, get_new_pid());
	return tsk;
}

static inline int init_task_new_context(struct task_desc * tsk, 
	unsigned long argv, unsigned long envp, unsigned long stack)
{
	/* 1. 初始化内核态上下文 */
	arch_init_kernel_regs(tsk, stack, argv, envp);
	/* 2. 初始化用户态上下文 */
	arch_init_user_regs(tsk, stack);
	return 0;
}

int init_task_sched(struct task_desc * tsk, int sched_policy)
{
	/**
	 * 根据调度策略决定调度类
	 */
	// switch (sched_policy)
	// {
	// case SCHED_FIFO:
	// case SCHED_TIMESLICE:
	// 	tsk->sched_class = &rt_sched_class;
	// 	break;
	// case SCHED_NORMAL:
	// case SCHED_SERVICE:
	// 	tsk->sched_class = &fair_sched_class;
	// 	break;
	// case SCHED_RMS:
	// 	tsk->sched_class = &rms_sched_class;
	// 	break;
	// default:
	// 	return -1;
	// }
	/**
	 * 所有调度数据都需要进行初始化，避免切换调度类时数据错误。
	 */
	tsk->ts_time = 0;

	tsk->sched_policy = sched_policy;
	return 0;
}

struct task_desc *create_task(struct task_create_param *param)
{
	unsigned long flags;
	struct task_desc *tsk;
	union process_union *stack = NULL;

	if (param->prio < 0 || param->prio >= MAX_RT_PRIO)
		goto out;

	if (!(tsk = create_task_desc()))
		goto out;

	if (!(stack = init_task_process_union(tsk)))
		goto out;

	if (init_task_new_context(tsk, (unsigned long)param->argv, 
		(unsigned long)param->envp, (unsigned long)stack))
		goto out;

	/* 利用 param 对 task 赋值 */
	tsk->sched_prio = param->prio;
	tsk->prio = param->prio;
	tsk->task_main = param->func;
	tsk->main_data = param->data;
	tsk->create_mode = param->app_flags;
	
	strncpy(tsk->name, param->name, TASK_NAME_LEN);
	strncpy(tsk->bin_path, param->name, BIN_PATH_MAX);

	/* 对调度进行初始化 */
	if (init_task_sched(tsk, param->sched_policy))
		goto out;

	/**
	 * 初始化退出信号值
	 */
	tsk->exit_signal = SIGCHLD & CSIGNAL;

	/* 初始化调度文件系统 */
	BUG_ON(init_task_fs(param->parent, tsk));

	pr_debug("%s : %s, %lu, %lx\n",
		__func__, tsk->name, tsk,  &task_entry);


	hold_task_desc(tsk);
	smp_lock_irqsave(&lock_all_task_list, flags);
	tsk->state = TASK_RUNNING;
	add_to_runqueue(tsk);
	list_insert_behind(&tsk->all_list, &sched_all_task_list);
	smp_unlock_irqrestore(&lock_all_task_list, flags);
	return tsk;
out:

	if (stack != NULL)
		free_proc_stack(stack);

	if (tsk->fp_regs)
		kfree(tsk->fp_regs);

	if (tsk)
		kfree(tsk);
	return NULL;
}

/**
 * 暂停任务运行
 * 在SMP下面应当重新实现本函数
 */
int suspend_task(struct task_desc *tsk)
{
	int ret = -1;
	unsigned long flags;

	if (!tsk)
		return -EINVAL;

	tsk->state |= TASK_STOPPED;
	smp_lock_irqsave(&lock_all_task_list, flags);
	del_from_runqueue(tsk);
	smp_unlock_irqrestore(&lock_all_task_list, flags);

	if (tsk == current)
		schedule();

	ret = 0;
	return ret;
}

/**
 * 恢复任务运行
 */
int resume_task(struct task_desc *tsk)
{
	int ret = -1;
	unsigned long flags;

	if (!tsk)
		return -EINVAL;

	smp_lock_irqsave(&lock_all_task_list, flags);
	tsk->state &= ~TASK_STOPPED;
	if (!tsk->state) {
		add_to_runqueue(tsk);
		tsk->state = TASK_RUNNING;
	}
	smp_unlock_irqrestore(&lock_all_task_list, flags);

	ret = 0;
	return ret;
}

void init_idle_process(union process_union *stack,
		struct task_desc *proc, int cpu)
{
	if (stack == NULL || proc == NULL)
		BUG();
	
	memset(proc, 0, sizeof(*proc));
	stack->process_desc.task = proc;
	stack->process_desc.cpu = cpu;
	proc->stack = stack->stack;

	snprintf(proc->name, TASK_NAME_LEN, "idle%d", cpu);
	proc->magic = TASK_MAGIC; 
	proc->state = TASK_INIT;
	proc->sched_prio = MAX_RT_PRIO;
	proc->prio = MAX_RT_PRIO;
	proc->flags = 0;
	proc->exit_code = 0;
	proc->task_main = &cpu_idle;
	proc->main_data = NULL;
	proc->in_run_queue = 0;
	proc->prev_sched = NULL;

	init_task_fs(NULL, proc);

	stack->process_desc.preempt_count = 0;

	list_init(&proc->run_list);
	list_init(&proc->all_list);
	init_waitqueue(&proc->wait_child_exit);
	hold_task_desc(proc);
}


/**
 * 调度初始化函数
 */
void init_sched(void)
{
	int i;

	list_init(&sched_all_task_list);
	
	if (!irqs_disabled()) {
		printk("O_O, disable irq, please!");
		local_irq_disable();
	}
	
	for(i = 0; i < ARRAY_SIZE(sched_runqueue_list); i++)
		list_init(&(sched_runqueue_list[i]));

	memset(sched_runqueue_mask, 0, sizeof(sched_runqueue_mask));
}
